
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain.embeddings.dashscope import DashScopeEmbeddings
from langchain_chroma import Chroma
from langchain.agents import AgentExecutor, create_tool_calling_agent, tool, \
  ZeroShotAgent
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain.tools.retriever import create_retriever_tool
from dotenv import find_dotenv, load_dotenv
import os

load_dotenv(find_dotenv())
DASHSCOPE_API_KEY = os.environ["DASHSCOPE_API_KEY"]

_context = {
  "context": "12313",
  "question": lambda x: x["standalone_question"],
}

_context2 = {
  "context": "123131111",
  "question": lambda x: x["standalone_question"],
}

class MyRetriever:
  def get_relevant_documents(self, query):
    print(query)
    return "123123"

class MyRetriever1(MyRetriever):
  def get_relevant_documents(self, query):
    print(query)
    print(self)
    return "123123"

_inputs = MyRetriever()

def build_chain():

  llm = ChatTongyi(
      model_name="qwen2-72b-instruct",
      streaming=True,
      api_key=os.environ["DASHSCOPE_API_KEY"]
  )
  embeddings = DashScopeEmbeddings(
      model="text-embedding-v1",
  )

  vector = Chroma(collection_name='customer', embedding_function=embeddings,
                  persist_directory='./chroma')
  retriever = vector.as_retriever(search_type="mmr",
                                  search_kwargs={'k': 6, 'lambda_mult': 0.25})

  template = """Answer the question based only on the following context:
    {context}

    Question: {question}
    """

  prompt = ChatPromptTemplate.from_template(template)

  chain = (
      {"context": retriever, "question": RunnablePassthrough()}
      | prompt
      | llm
      | StrOutputParser()
  )
  return chain


if __name__ == '__main__':
  # chain = build_chain()
  # print(type(chain))
  # aa = build_chain().invoke("RAG的本质是什么？")
  # print(aa)
  cha = (
      1 | 2
  )
  print(cha)